import acl
import sys
import struct
import numpy as np
sys.path.append("..")
from AscendResource.AscendResource import ascend_resource
from CommonModule.LogModule import DemoLogger, check_ret
from CommonModule.ConstModule import ACL_MEMCPY_HOST_TO_HOST, \
                     ACL_MEMCPY_HOST_TO_DEVICE, \
                     ACL_MEMCPY_DEVICE_TO_HOST, \
                     ACL_MEMCPY_DEVICE_TO_DEVICE, \
                     ACL_MEM_MALLOC_HUGE_FIRST, \
                     ACL_MEM_MALLOC_HUGE_ONLY, \
                     ACL_MEM_MALLOC_NORMAL_ONLY, \
                     PIXEL_FORMAT_YUV_SEMIPLANAR_420, \
                     PIXEL_FORMAT_YVU_SEMIPLANAR_420, \
                     PIXEL_FORMAT_YUV_SEMIPLANAR_422, \
                     PIXEL_FORMAT_YVU_SEMIPLANAR_422, \
                     PIXEL_FORMAT_YUV_SEMIPLANAR_444, \
                     PIXEL_FORMAT_YVU_SEMIPLANAR_444, \
                     DVPP, DEVICE, HOST, \
                     DEBUG,INFO,WARNING,ERROR, \
                     STATIC_MODE, DYNAMIC_BATCH, DYNAMIC_RESOLUTION, \
                     NPY_BYTE

class AscendModel():
    def __init__(self,
                 context,
                 stream,
                 model_path):
        self.class_name = self.__class__.__name__
        self.context = context  # pointer
        self.stream = stream

        self.model_desc = None  # pointer when using
        self.model_path = model_path  # string
        self.model_id = None  # pointer
        self.model_input_ptr = None
        self.model_output_ptr = None
        self.model_input_size = None
        self.model_output_size = None
        self.model_batch = None
        self.input_dataset = None
        self.output_dataset = None
        self.host_output = []
        self.output_info = []
   
        self.memory = ascend_resource.memory_pool       

        self.init_resource()
        self.is_create = True

    def init_resource(self):
        #acl.rt.set_context(self.context)
        DemoLogger.do_log(INFO, self.class_name, 'init_resource', ' init resource stage start')
        self.gen_model_desc()
        self.init_dataset()
        DemoLogger.do_log(INFO, self.class_name, 'init_resource', ' init resource stage success')

    def init_dataset(self):
        self.input_dataset = acl.mdl.create_dataset()
        self.output_dataset = acl.mdl.create_dataset()
    
    def gen_model_desc(self):
        self.model_id, ret = acl.mdl.load_from_file(self.model_path)
        check_ret(self.class_name, '_gen_model_desc', 'fail to load mode from file', ret)
        self.model_desc = acl.mdl.create_desc()
        ret = acl.mdl.get_desc(self.model_desc, self.model_id)
        check_ret(self.class_name, '_gen_model_desc', 'fail to get model id', ret)    

    def gen_model_data_num(self, opt):
        if opt == 'input':
            size = acl.mdl.get_num_inputs(self.model_desc)
        elif opt == 'output':
            size = acl.mdl.get_num_outputs(self.model_desc)
        else:
            DemoLogger.do_log(ERROR, self.class_name, 'gen_model_data_num', 'opt is not supported')
        return size

    def gen_datainfo(self,opt):
        size = self.gen_model_data_num(opt)
        ptr_list = [0] * size
        size_list = [0] * size
        batch_list = [0]* size
        for i in range(0, size):
            if opt == 'input':
                size_list[i] = acl.mdl.get_input_size_by_index(self.model_desc, i)
                ptr_list[i] = self.memory.memory_allocation_(size_list[i], DEVICE)
                dims, ret = acl.mdl.get_input_dims(self.model_desc,i)
                check_ret(self.class_name, 'send_data_to_model', 'fail to get_input_dims', ret)
                batch_list[i] = dims["dims"][0]
            if opt == 'output':
                size_list[i] = acl.mdl.get_output_size_by_index(self.model_desc, i)
                ptr_list[i] = self.memory.memory_allocation_(size_list[i], DEVICE)
                dims, ret = acl.mdl.get_output_dims(self.model_desc,i)
                check_ret(self.class_name, 'send_data_to_model', 'fail to get_output_dims', ret)
                batch_list[i] = dims["dims"][0]
                datatype = acl.mdl.get_output_data_type(self.model_desc, i)
                self.output_info.append({"shape": tuple(dims["dims"]), "type": datatype})

        DemoLogger.do_log(INFO, self.class_name, 'gen_model_data', 'init model {} success'.format(opt))
        return ptr_list, size_list, batch_list 
 
    def gen_databuffer(self, opt, buffer_ptr:list, buffer_size:list):
        if opt == 'input':
            dataset = self.input_dataset
        if opt == 'output':
            dataset = self.output_dataset

        for i in range(0, len(buffer_ptr)):
            temp_buffer = acl.create_data_buffer(buffer_ptr[i], buffer_size[i])
            _, ret = acl.mdl.add_dataset_buffer(dataset, temp_buffer)
            try :
                check_ret(self.class_name, '_gen_databuffer', 'fail to create {} dataset buffer'.format(opt),ret)
            except :
                acl.destroy_data_buffer(dataset)
                return
        DemoLogger.do_log(INFO,self.class_name, '_gen_databuffer', 'create model data buffer success')

    def gen_dataset(self, policy, width, height, batch):
        if policy == STATIC_MODE:
           self.policy = policy
           if self.is_create == True:
                self.is_create = False
                #input
                self.model_input_ptr, self.model_input_size, self.model_batch = self.gen_datainfo('input')
                self.gen_databuffer('input', self.model_input_ptr, self.model_input_size)
                #output
                self.model_output_ptr, self.model_output_size, _ = self.gen_datainfo('output')
                self.gen_databuffer('output', self.model_output_ptr, self.model_output_size)
        elif policy == DYNAMIC_BATCH:
            pass
        elif policy == DYNAMIC_RESOLUTION:
            pass
        else:
            DemoLogger.do_log(ERROR,self.class_name, 'gen_dataset', 'policy is not supported')

    def run(self, policy, dvpp_buffer: list, dvpp_size: list, width = None, height = None, batch = None):
        self.gen_dataset(policy, width, height, batch)
        self.send_data_to_model(dvpp_buffer, dvpp_size)      
        self.forward()
        self.send_result_to_host()
        return self.host_output

    def forward(self):
        ret = acl.mdl.execute_async(self.model_id,
                              self.input_dataset,
                              self.output_dataset,
                              self.stream)
        check_ret(self.class_name, 'forward', 'fail to send desc to model', ret)
        ret = acl.rt.synchronize_stream(self.stream)
        check_ret(self.class_name, 'forward', 'fail to execute model', ret)
        DemoLogger.do_log(INFO, self.class_name, 'forward', 'model execute stage success')

    def send_data_to_model(self, ptr:list, size:list):
        num = self.gen_model_data_num('input')
        if len(ptr) != len(size) or num != len(ptr):
           DemoLogger.do_log(ERROR, self.class_name, 'send_data_to_model', 'input data num do not match model requirement')
        for i in range(0, num):
            if size[i] == self.model_input_size[i]:                   
                ret = acl.rt.memcpy(self.model_input_ptr[i], size[i], ptr[i], size[i], ACL_MEMCPY_DEVICE_TO_DEVICE)
                check_ret(self.class_name, 'send_data_to_model', 'fail to memcpy for ptr[{}]'.format(i), ret)
            else:
                DemoLogger.do_log(ERROR, self.class_name, 'send_data_to_model', 'input data size do not match model requirement')
        DemoLogger.do_log(INFO, self.class_name, 'send_data_to_model', 'send data to model success')

    def send_result_to_host(self):
        self.host_output = []
        num = acl.mdl.get_dataset_num_buffers(self.output_dataset)
        for i in range(num):
            temp_output_buffer = acl.mdl.get_dataset_buffer(self.output_dataset, i)
            infer_output_ptr = acl.get_data_buffer_addr(temp_output_buffer)
            infer_output_size = acl.get_data_buffer_size(temp_output_buffer)
            output_host = self.memory.memory_allocation_(infer_output_size, HOST)
            ret = acl.rt.memcpy(output_host, infer_output_size, infer_output_ptr,
                                infer_output_size, ACL_MEMCPY_DEVICE_TO_HOST)
            check_ret(self.class_name, 'send_result_to_host', 'fail to memcy to host', ret)
            self.host_output.append([output_host, infer_output_size])
        DemoLogger.do_log(INFO, self.class_name, 'send_result_to_host', 'send data to host success')

    def _release_databuffer(self,opt):
        if opt == 'input':
            dataset = self.input_dataset
        if opt == 'output':
            dataset = self.output_dataset

        for i in range(0, acl.mdl.get_dataset_num_buffers(dataset)):
            data_buf = acl.mdl.get_dataset_buffer(dataset, i)
            if data_buf:
                ret = acl.destroy_data_buffer(data_buf)
                check_ret(self.class_name, '_release_dataset', 'fail to destory {} databuffer'.format(opt), ret)

    def _release_dataset(self):
        if self.input_dataset:
            ret = acl.mdl.destroy_dataset(self.input_dataset)
            check_ret(self.class_name, '_release_dataset', 'fail to destory input dataset', ret)
        if self.output_dataset:
            ret = acl.mdl.destroy_dataset(self.output_dataset)
            check_ret(self.class_name, '_release_dataset', 'fail to destory output dataset', ret)

    def release_model_resource(self):
        self._release_databuffer('input')
        self._release_databuffer('output')
        self._release_dataset()
        if self.model_id:
            ret = acl.mdl.unload(self.model_id)
            check_ret(self.class_name, 'release_model_resource', 'fail to unload mode', ret)
        if self.model_desc:
            ret = acl.mdl.destroy_desc(self.model_desc)
            check_ret(self.class_name, 'release_model_resource', 'fail to destory desc', ret)
        DemoLogger.do_log(INFO, self.class_name, 'release_model_resource','release source success')

    def __del__(self):
        self.release_model_resource()
    
